import rospy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import os

"""
数据集制作工具
根据所提供的话题和ros bag，按照一定的频率截取图像用于制作数据集
"""

class ImageSaver:
    def __init__(self):
        self.bridge = CvBridge()
        self.frame_count = 0

        ############需要修改###################
        self.image_count = 337  # 生成图像的起始编号，注意要改成当前文件夹里最后一个图像的标号+1，不然会覆盖
        self.save_interval = 5  # 每隔五帧保存一次图像
        self.image_topic = "/perception/zed2i/rgb_leftView"     # 替换成你的图像话题
        self.save_folder = "/home/hw/dataset/my_dataset2/images" # 替换成你的保存文件夹路径
        #####################################
        
        if not os.path.exists(self.save_folder):
            os.makedirs(self.save_folder)
        self.image_sub = rospy.Subscriber(self.image_topic, Image, self.image_callback)

    def image_callback(self, msg):
        if self.frame_count % self.save_interval == 0:
            try:
                cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
                image_name = f"{self.image_count}.jpg"
                image_path = os.path.join(self.save_folder, image_name)
                cv2.imwrite(image_path, cv_image)
                print(f"Saved image: {image_path}")
                self.image_count += 1
            except Exception as e:
                print(f"Error processing image: {e}")

        self.frame_count += 1

if __name__ == "__main__":
    rospy.init_node('image_saver', anonymous=True)
    image_saver = ImageSaver()
    try:
        rospy.spin()
    except KeyboardInterrupt:
        print("Shutting down")
    cv2.destroyAllWindows()